"""
This module provides analysis of pulsed measurements.

The first part provides simple numeric functions.

The second part provides Trait GUIs
"""

import numpy as np
from fitting import find_edge

def spin_state(c, dt, T, t0=0.0, t1=-1.):
    
    """
    Compute the spin state from a 2D array of count data.
    
    Parameters:
    
        c    = count data
        dt   = time step
        t0   = beginning of integration window relative to the edge
        t1   = None or beginning of integration window for normalization relative to edge
        T    = width of integration window
        
    Returns:
    
        y       = 1D array that contains the spin state
        profile = 1D array that contains the pulse profile
        edge    = position of the edge that was found from the pulse profile
        
    If t1<0, no normalization is performed. If t1>=0, each data point is divided by
    the value from the second integration window and multiplied with the mean of
    all normalization windows.
    """

    profile = c.sum(0)
    edge = find_edge(profile)
    
    I = int(round(T/float(dt)))
    i0 = edge + int(round(t0/float(dt)))
    y = np.empty((c.shape[0],))
    for i, slot in enumerate(c):
        y[i] = slot[i0:i0+I].sum()
    if t1 >= 0:
        i1 = edge + int(round(t1/float(dt)))    
        y1 = np.empty((c.shape[0],))
        for i, slot in enumerate(c):
            y1[i] = slot[i1:i1+I].sum()
        y = y/y1*y1.mean()
    return y, profile, edge


#########################################
# Trait GUIs for pulsed fits
#########################################

from traits.api import Instance, Any, Property, Range, Float, Int, Bool, Array, List, Str, Tuple, Enum,\
                                 on_trait_change, cached_property, DelegatesTo
from traitsui.api import View, Item, Tabbed, Group, HGroup, VGroup, VSplit, EnumEditor, TextEditor, InstanceEditor
from enable.api import ComponentEditor
from chaco.api import ArrayDataSource, LinePlot, LinearMapper, ArrayPlotData, Spectral, PlotLabel

from tools.chaco_addons import SavePlot as Plot, SaveTool

import threading
import time
import logging

import fitting

from tools.utility import GetSetItemsMixin

from measurements.pulsed import Pulsed, PulsedTau

class PulsedTool( GetSetItemsMixin ):

    """
    Base class for a pulsed analysis. Provides calculation of spin state
    and plotting.
    Derive from this to create analysis tools for pulsed measurements.
    """

    # the measurement to analyze
    measurement = Any(editor=InstanceEditor)
    
    # plotting
    matrix_data = Instance( ArrayPlotData)
    line_data   = Instance( ArrayPlotData )
    pulse_data  = Instance( ArrayPlotData )
    matrix_plot = Instance( Plot, editor=ComponentEditor() )
    pulse_plot  = Instance( Plot, editor=ComponentEditor() )
    line_plot   = Instance( Plot, editor=ComponentEditor() )

    get_set_items = ['__doc__','measurement']

    traits_view = View(VGroup(Item(name='measurement', style='custom', show_label=False),
                              VSplit(Item('matrix_plot', show_label=False, width=500, height=-300, resizable=True),
                                     Item('line_plot', show_label=False, width=500, height=-300, resizable=True),
                                     Item('pulse_plot', show_label=False, width=500, height=-300, resizable=True),
                                     ),
                              ),
                       title='Pulsed Tool',
                       buttons=[],
                       resizable=True,
                       height=-640
                       )

    def __init__(self, **kwargs):
        super(PulsedTool, self).__init__(**kwargs)
        self._create_matrix_plot()
        self._create_pulse_plot()
        self._create_line_plot()
        self.on_trait_change(self._update_matrix_index, 'measurement.time_bins,measurement.n_laser', dispatch='ui')
        self.on_trait_change(self._update_matrix_value, 'measurement.count_data', dispatch='ui')
        self.on_trait_change(self._update_pulse_index, 'measurement.time_bins', dispatch='ui')
        self.on_trait_change(self._update_pulse_value, 'measurement.pulse', dispatch='ui')
        self.on_trait_change(self._update_line_plot_value, 'measurement.spin_state', dispatch='ui')
        self.on_trait_change(self._on_edge_change, 'measurement.edge', dispatch='ui')

    # plotting
    def _create_matrix_plot(self):
        matrix_data = ArrayPlotData(image=np.zeros((2,2)))
        plot = Plot(matrix_data, width=500, height=500, resizable='hv', padding=8, padding_left=48, padding_bottom=36)
        plot.index_axis.title = 'time [ns]'
        plot.value_axis.title = 'laser pulse #'
        plot.img_plot('image',
                      xbounds=(0,1),
                      ybounds=(0,1),
                      colormap=Spectral)[0]
        plot.tools.append(SaveTool(plot))
        self.matrix_data = matrix_data
        self.matrix_plot = plot
    
    def _create_pulse_plot(self):
        pulse_data  = ArrayPlotData(x=np.array((0.,0.1,0.2)),y=np.array((0,1,2)))
        plot = Plot(pulse_data, padding=8, padding_left=64, padding_bottom=36)    
        line = plot.plot(('x','y'), style='line', color='blue', name='data')[0]
        plot.index_axis.title = 'time [ns]'
        plot.value_axis.title = 'intensity'
        edge_marker = LinePlot(index = ArrayDataSource(np.array((0,0))),
                               value = ArrayDataSource(np.array((0,1e9))),
                               color = 'red',
                               index_mapper = LinearMapper(range=plot.index_range),
                               value_mapper = LinearMapper(range=plot.value_range),
                               name='marker')
        plot.add(edge_marker)
        plot.tools.append(SaveTool(plot))
        self.pulse_data = pulse_data
        self.pulse_plot = plot
        
    def _create_line_plot(self):
        line_data   = ArrayPlotData(index=np.array((0,1)), spin_state=np.array((0,0)),)
        plot = Plot(line_data, padding=8, padding_left=64, padding_bottom=36)
        plot.plot(('index','spin_state'), color='blue', name='spin_state')
        plot.index_axis.title = 'pulse #'
        plot.value_axis.title = 'spin state'
        plot.tools.append(SaveTool(plot))
        self.line_data = line_data
        self.line_plot = plot

    def _update_matrix_index(self):
        if self.measurement is not None:
            self.matrix_plot.components[0].index.set_data((self.measurement.time_bins[0], self.measurement.time_bins[-1]),(0.0,float(self.measurement.n_laser)))
        
    def _update_matrix_value(self):
        if self.measurement is not None:
            s = self.measurement.count_data.shape
            if not s[0]*s[1] > 1000000:
                self.matrix_data.set_data('image', self.measurement.count_data)
    def _update_pulse_index(self):
        if self.measurement is not None:
            self.pulse_data.set_data('x', self.measurement.time_bins)        
    def _update_pulse_value(self):
        if self.measurement is not None:
            self.pulse_data.set_data('y', self.measurement.pulse)
    def _on_edge_change(self):
        if self.measurement is not None:
            y=self.measurement.edge
            self.pulse_plot.components[1].index.set_data(np.array((y,y)))
    def _update_line_plot_value(self):
        if self.measurement is not None:
            y=self.measurement.spin_state
            n = len(y)
            old_index = self.line_data.get_data('index')
            if old_index is not None and len(old_index) != n:
                self.line_data.set_data('index',np.arange(n))
            self.line_data.set_data('spin_state',y)

    def save_matrix_plot(self, filename):
        save_figure(self.matrix_plot, filename)
    
    def save_line_plot(self, filename):
        save_figure(self.line_plot, filename)


class PulsedToolTau( PulsedTool ):

    """
    Analysis of a pulsed measurement with a 'tau' as index-data.
    """

    # overwrite __init__ such that change of 'tau' causes plot update 
    def __init__(self, **kwargs):
        super(PulsedToolTau, self).__init__(**kwargs)
        self.on_trait_change(self._on_tau_change, 'measurement.tau', dispatch='ui')
        #self.on_trait_change(self._on_tau_change, 'measurement.tau') # ToDo: fix this

    # overwrite the line_plot such that the x-axis label is time 
    def _create_line_plot(self):
        line_data   = ArrayPlotData(index=np.array((0,1)), spin_state=np.array((0,0)),)
        plot = Plot(line_data, padding=8, padding_left=64, padding_bottom=36)
        plot.plot(('index','spin_state'), color='blue', name='spin_state')
        plot.index_axis.title = 'time [micro s]'
        plot.value_axis.title = 'spin state'
        plot.tools.append(SaveTool(plot))
        self.line_data = line_data
        self.line_plot = plot

    # overwrite this one to throw out setting of index data according to length of spin_state
    def _update_line_plot_value(self):
        y = self.measurement.spin_state
        self.line_data.set_data('spin_state',y)

    # provide method for update of tau
    def _on_tau_change(self):
        self.line_data.set_data('index',self.measurement.tau*1e-3)

    # overwrite this to change the window title
    traits_view = View(VGroup(Item(name='measurement', style='custom', show_label=False),
                              VSplit(Item('matrix_plot', show_label=False, width=500, height=300, resizable=True),
                                     Item('line_plot', show_label=False, width=500, height=300, resizable=True),
                                     Item('pulse_plot', show_label=False, width=500, height=300, resizable=True),
                                     ),
                              ),
                       title='Pulsed Analysis Tau',
                       buttons=[], resizable=True
                       )


class FitToolTau( PulsedToolTau ):

    """
    Base class for PulsedTool with a tau and fit.
    """

    # fit results
    fit_result = Tuple()
    label_text = Str('')

    # add fit results to the get_set_items
    get_set_items = PulsedToolTau.get_set_items + ['fit_result','label_text']

    # overwrite __init__ to trigger update events
    def __init__(self, **kwargs):
        super(FitToolTau, self).__init__(**kwargs)
        self.on_trait_change(self._update_fit, 'measurement.spin_state', dispatch='ui')
        self.on_trait_change(self._on_fit_result_change, 'fit_result', dispatch='ui')
        self.on_trait_change(self._on_label_text_change, 'label_text', dispatch='ui')
    
    def _update_fit(self):
        pass
        
    # overwrite the line_plot to include fit and text label 
    def _create_line_plot(self):
        line_data   = ArrayPlotData(index=np.array((0,1)),
                                    spin_state=np.array((0,0)),
                                    fit=np.array((0,0)))
        plot = Plot(line_data, padding=8, padding_left=64, padding_bottom=36)
        plot.plot(('index','spin_state'), color='blue', name='spin_state')
        plot.plot(('index','fit'), color='red', name='fit')
        plot.index_axis.title = 'time [micro s]'
        plot.value_axis.title = 'spin state'
        plot.overlays.insert(0, PlotLabel(text=self.label_text, hjustify='left', vjustify='bottom', position=[64,32]) )
        plot.tools.append(SaveTool(plot))
        self.line_data = line_data
        self.line_plot = plot

    def _on_fit_result_change(self, new):
        pass
    
    def _on_label_text_change(self, new):
        self.line_plot.overlays[0].text = new    
    



#########################################
# testing
#########################################

if __name__ == '__main__':

    logging.getLogger().addHandler(logging.StreamHandler())
    logging.getLogger().setLevel(logging.DEBUG)
    logging.getLogger().info('Starting logger.')
    
    from tools.emod import JobManager
    
    JobManager().start()
        
    import measurements.pulsed
    
    # basic example to analyze a Rabi measurement
    #m = measurements.pulsed.Rabi()
    #a = PulsedTool()
    #a = PulsedToolTau()
    #a = PulsedToolDoubleTauRef()
    a = RabiAna()
    #a = PulseCalAna()
    #a.measurement = measurements.pulsed.Rabi()
    #a.measurement = mp.PulseCal()
    a.edit_traits()
    #a.load('')
    #a.measurement = m
